package com.media.codec;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import lombok.Data;

import java.util.ArrayList;
import java.util.List;

/**
 * 0                   1                   2                   3
 * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |V=2|P|X|  CC   |M|     PT      |       sequence number         |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |                           timestamp                           |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |           synchronization source (SSRC) identifier            |
 * +=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
 * |            contributing source (CSRC) identifiers             |
 * |                             ....                              |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * <p>
 * 0                   1                   2                   3
 * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |      defined by profile       |           length              |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |                        header extension                       |
 * |                             ....                              |
 *
 */
@Data
public class RtpPacket {

    // internal vars --------------------------------------------------------------------------------------------------

    private RtpVersion version;
    private boolean marker;
    private int payloadType;
    private int sequenceNumber;
    private long timestamp;
    private long ssrc;

    private short extensionHeaderData;
    private byte[] extensionData;

    private List<Long> contributingSourceIds;

    private ByteBuf data;

    // constructors ---------------------------------------------------------------------------------------------------

    public RtpPacket() {
        this.version = RtpVersion.V2;
    }

    // public static methods ------------------------------------------------------------------------------------------

    public static RtpPacket decode(byte[] data) {
        return decode(Unpooled.wrappedBuffer(data));
    }

    public static RtpPacket decode(ByteBuf buffer) throws IndexOutOfBoundsException {
        if (buffer.readableBytes() < 12) {
            throw new IllegalArgumentException("A RTP packet must be at least 12 octets long");
        }
        RtpPacket packet = new RtpPacket();

        // Version, Padding, eXtension, CSRC Count
        byte b = buffer.readByte();
        RtpVersion version = RtpVersion.fromByte(b);
        if (version != RtpVersion.V0 && version != RtpVersion.V1 && version != RtpVersion.V2) {
            return null;
        }
        packet.version = version;
        boolean padding = (b & 0x20) > 0; // mask 0010 0000
        boolean extension = (b & 0x10) > 0; // mask 0001 0000
        int contributingSourcesCount = b & 0x0f; // mask 0000 1111

        // Marker, Payload Type
        b = buffer.readByte();
        packet.marker = (b & 0x80) > 0; // mask 0000 0001
        packet.payloadType = (b & 0x7f); // mask 0111 1111

        packet.sequenceNumber = buffer.readUnsignedShort();
        packet.timestamp = buffer.readUnsignedInt();
        packet.ssrc = buffer.readUnsignedInt();

        // Read extension headers & data
        if (extension) {
            packet.extensionHeaderData = buffer.readShort();
            packet.extensionData = new byte[buffer.readUnsignedShort()];
            buffer.readBytes(packet.extensionData);
        }

        // Read CCRC's
        if (contributingSourcesCount > 0) {
            packet.contributingSourceIds = new ArrayList<>(contributingSourcesCount);
            for (int i = 0; i < contributingSourcesCount; i++) {
                long contributingSource = buffer.readUnsignedInt();
                packet.contributingSourceIds.add(contributingSource);
            }
        }

        if (!padding) {
            // No padding used, assume remaining data is the packet
            byte[] remainingBytes = new byte[buffer.readableBytes()];
            buffer.readBytes(remainingBytes);
            packet.setData(remainingBytes);
        } else {
            // Padding bit was set, so last byte contains the number of padding octets that should be discarded.
            short lastByte = buffer.getUnsignedByte(buffer.readerIndex() + buffer.readableBytes() - 1);
            byte[] dataBytes = new byte[buffer.readableBytes() - lastByte];
            buffer.readBytes(dataBytes);
            packet.setData(dataBytes);
            // Discard rest of buffer.
            buffer.skipBytes(buffer.readableBytes());
        }

        return packet;
    }

    public static ByteBuf encode(int fixedBlockSize, RtpPacket packet) {
        int size = 12; // Fixed width
        if (packet.hasExtension()) {
            size += 4 + packet.getExtensionDataSize();
        }
        size += packet.getContributingSourcesCount() * 4;
        size += packet.getDataSize();

        // If packet was configured to have padding (fixed block size), calculate padding and add it.
        int padding = 0;
        if (fixedBlockSize > 0) {
            // If padding modulus is > 0 then the padding is equal to:
            // (global size of the compound RTCP packet) mod (block size)
            // Block size alignment might be necessary for some encryption algorithms
            // RFC section 6.4.1
            padding = fixedBlockSize - (size % fixedBlockSize);
            if (padding == fixedBlockSize) {
                padding = 0;
            }
        }

        ByteBuf headBuf = Unpooled.directBuffer(12);
        // Version, Padding, eXtension, CSRC Count
        byte b = packet.getVersion().getByte();
        if (padding > 0) {
            b |= 0x20;
        }
        if (packet.hasExtension()) {
            b |= 0x10;
        }
        b |= packet.getContributingSourcesCount();
        headBuf.writeByte(b);

        // Marker, Payload Type
        b = (byte) packet.getPayloadType();
        if (packet.getMarker()) {
            b |= 0x80; // 1000 0000
        }
        headBuf.writeByte(b);

        headBuf.writeShort(packet.sequenceNumber);
        headBuf.writeInt((int) packet.timestamp);
        headBuf.writeInt((int) packet.ssrc);

        // Write extension headers & data
        if (packet.hasExtension()) {
            headBuf.writeShort(packet.extensionHeaderData);
            headBuf.writeShort(packet.extensionData.length);
            headBuf.writeBytes(packet.extensionData);
        }

        // Write CCRC's
        if (packet.getContributingSourcesCount() > 0) {
            for (Long contributingSourceId : packet.getContributingSourceIds()) {
                headBuf.writeInt(contributingSourceId.intValue());
            }
        }

        if (padding > 0) {
            // Final bytes: padding
            ByteBuf paddingBuf = Unpooled.directBuffer(8 * padding);
            for (int i = 0; i < (padding - 1); i++) {
                paddingBuf.writeByte(0x00);
            }

            // Final byte: the amount of padding bytes that should be discarded.
            // Unless something's wrong, it will be a multiple of 4.
            paddingBuf.writeByte(padding);

            return Unpooled.wrappedBuffer(headBuf, packet.data, paddingBuf);
        }
        return Unpooled.wrappedBuffer(headBuf, packet.data);
    }

    // public methods ----------------------------------------------------------------------------------------------

    public ByteBuf encode(int fixedBlockSize) {
        return encode(fixedBlockSize, this);
    }

    public ByteBuf encode() {
        return encode(0, this);
    }

    public void addContributingSourceId(long contributingSourceId) {
        if (this.contributingSourceIds == null) {
            this.contributingSourceIds = new ArrayList<>();
        }
        this.contributingSourceIds.add(contributingSourceId);
    }

    public int getDataSize() {
        if (this.data == null) {
            return 0;
        }
        return this.data.capacity();
    }

    public int getExtensionDataSize() {
        if (this.extensionData == null) {
            return 0;
        }
        return this.extensionData.length;
    }

    public int getContributingSourcesCount() {
        if (this.contributingSourceIds == null) {
            return 0;
        }
        return this.contributingSourceIds.size();
    }

    public void setExtensionHeader(short extensionHeaderData, byte[] extensionData) {
        if (extensionData.length > 65536) {
            throw new IllegalArgumentException("Extension data cannot exceed 65536 bytes");
        }
        this.extensionHeaderData = extensionHeaderData;
        this.extensionData = extensionData;
    }

    // getters & setters ----------------------------------------------------------------------------------------------

    public void setVersion(RtpVersion version) {
        if (version != RtpVersion.V2) {
            throw new IllegalArgumentException("Only V2 is supported");
        }
        this.version = version;
    }

    public boolean getMarker() {
        return marker;
    }

    public boolean hasExtension() {
        return this.extensionData != null;
    }

    public void setPayloadType(int payloadType) {
        if ((payloadType < 0) || (payloadType > 127)) {
            throw new IllegalArgumentException("PayloadType must be in range [0;127]");
        }
        this.payloadType = payloadType;
    }

    public void setSsrc(long ssrc) {
        if ((ssrc < 0) || (ssrc > 0xffffffffL)) {
            throw new IllegalArgumentException("Valid range for SSRC is [0;0xffffffff]");
        }
        this.ssrc = ssrc;
    }

    public void setData(byte[] data) {
        this.data = Unpooled.wrappedBuffer(data);
    }

    public void setData(ByteBuf data) {
        this.data = data;
    }

    public byte[] getDataAsArray() {
        return null == this.data ? null : this.data.array();
    }

    // low level overrides --------------------------------------------------------------------------------------------

    @Override
    public String toString() {
        return "RtpPacket{V=" + this.version +
                ", X=" + this.hasExtension() +
                ", CC=" + this.getContributingSourcesCount() +
                ", M=" + this.marker +
                ", PT=" + this.payloadType +
                ", SN=" + this.sequenceNumber +
                ", TS=" + this.timestamp +
                ", SSRC=" + this.ssrc +
                ", CSRCs=" + this.contributingSourceIds +
                ", data=" + this.getDataSize() + " bytes}";
    }
}
